#!/usr/bin/python
# -*- coding: UTF-8 -*-

################################################################################
#
# Copyright (c) 2021 openEuler.org, Inc. All Rights Reserved
#
################################################################################
"""
Supplemental cve information

Authors: zhangjianjun
Date:    1/19/2021 11:01 AM
"""

from tabletask import crawltask
from dbConnecttion.MysqlConn import Mysql
import time
import datetime



def query_cve_all_data(mysql, days):
    """
    Find cve with missing necessary fields
    """
    before_date = (datetime.date.today() -
                   datetime.timedelta(days=days)).strftime("%Y-%m-%d %H:%M:%S")
    score_sql = "select cve_id,cve_num from cve_score " \
                "where nvd_score = 0 and update_time >= %s order by update_time desc"
    val = (before_date,)
    cve_list = []
    cve_id_list = []
    score_result = mysql.getMany(score_sql, val)
    if score_result and len(score_result) > 0:
        for sc in score_result:
            if sc and sc not in cve_id_list:
                cve_id_list.append(sc)

    cve_desc_sql = "select cve_id,cve_num from " \
                   "cve_vuln_center where (cve_desc = %s or " \
                   "repair_time = %s) and update_time >= %s " \
                   "order by cve_id desc"
    center_val = ("", "", before_date)
    center_result = mysql.getMany(cve_desc_sql, center_val)
    if center_result and len(center_result) > 0:
        for ce in center_result:
            if ce and ce not in cve_id_list:
                cve_id_list.append(ce)

    cve_template_sql = "select cve_id,cve_num from " \
                       "cve_issue_template where (nvd_score = %s or " \
                       "nvd_vector = %s or cve_brief = %s) and update_time >= %s " \
                       "order by cve_id desc"
    template_val = (0, "", "", before_date)
    template_result = mysql.getMany(cve_template_sql, template_val)
    if template_result and len(template_result) > 0:
        for ct in template_result:
            if ct and ct not in cve_id_list:
                cve_id_list.append(ct)

    if cve_id_list and len(cve_id_list) > 0:
        for ci in cve_id_list:
            center_sql = "select cve_num, pack_name, cve_version," \
                         "cve_desc,repair_time,cve_status,cve_id from " \
                         "cve_vuln_center where cve_id = %s and " \
                         "cve_num = %s order by cve_id desc"
            center_val = (ci["cve_id"], ci["cve_num"])
            center_result = mysql.getOne(center_sql, center_val)
            if center_result:
                cve_list.append(center_result)
    return cve_list


def update_cve(url, result_dict, cve, mysql):
    """
    update data
    """
    cve_num = str(cve["cve_num"]).strip()
    cve_version = str(cve["cve_version"]).strip()
    pack_name = str(cve["pack_name"]).strip()
    # Crawler web data
    crawl_list = crawltask.crawling(url)
    print(crawl_list)
    # Determine whether the database content is the latest data
    if crawl_list[0]:
        if str(result_dict["nvd_score"]) == str(crawl_list[0]) \
                and str(result_dict["vector_value"]) == str(crawl_list[4]):
            if result_dict['cve_status'] in [3, 4, 7]:
                print("update data:" + cve_num)
                update_time = str(time.strftime("%Y-%m-%d %H:%M:%S", time.localtime()))
                try:
                    sql = "update cve_origin_excel set nvd_score=%s, cve_level=%s, cve_desc=%s, " \
                          "repair_time=%s, vector_value=%s, attack_vector=%s, access_vector=%s, " \
                          "attack_complexity=%s, access_complexity=%s, privilege_required=%s, " \
                          "user_interaction=%s, scope=%s, confidentiality=%s, integrity=%s, " \
                          "availability=%s, authentication=%s, cve_status=%s, update_time=%s " \
                          "where cve_num=%s and pack_name = %s and cve_version = %s"
                    val = (
                        crawl_list[0], crawl_list[1], crawl_list[2], crawl_list[3], crawl_list[4],
                        crawl_list[5],
                        crawl_list[6], crawl_list[7], crawl_list[8], crawl_list[9],
                        crawl_list[10], crawl_list[11], crawl_list[12], crawl_list[13], crawl_list[14],
                        crawl_list[15], 1, update_time, cve_num, pack_name, cve_version)
                    mysql.update(sql, val)
                    mysql.dispose()
                except IndexError as e:
                    mysql.dispose(0)
                    print("Subscript out of bounds", e)
            else:
                print("The database is the latest data:" + cve_num)
        else:
            print("update data:" + cve_num)
            update_time = str(time.strftime("%Y-%m-%d %H:%M:%S", time.localtime()))
            try:
                sql = "update cve_origin_excel set nvd_score=%s, cve_level=%s, cve_desc=%s, " \
                      "repair_time=%s, vector_value=%s, attack_vector=%s, access_vector=%s, " \
                      "attack_complexity=%s, access_complexity=%s, privilege_required=%s, " \
                      "user_interaction=%s, scope=%s, confidentiality=%s, integrity=%s, " \
                      "availability=%s, authentication=%s, cve_status=%s, update_time=%s, " \
                      "score_type=%s " \
                      "where cve_num=%s and pack_name = %s and cve_version = %s"
                val = (
                    crawl_list[0], crawl_list[1], crawl_list[2], crawl_list[3], crawl_list[4],
                    crawl_list[5], crawl_list[6], crawl_list[7], crawl_list[8], crawl_list[9],
                    crawl_list[10], crawl_list[11], crawl_list[12], crawl_list[13], crawl_list[14],
                    crawl_list[15], 1, update_time, crawl_list[16],
                    cve_num, pack_name, cve_version)
                mysql.update(sql, val)
                mysql.dispose()
            except IndexError as e:
                print("Subscript out of bounds", e)
                mysql.dispose(0)
    else:
        print("error: ", result_dict)


def update_cve_vuln(url, cve, mysql):
    """
    update data
    """
    ret_value = 0
    cve_num = str(cve["cve_num"]).strip()
    cve_id = cve["cve_id"]
    cve_status = cve["cve_status"]
    if cve["cve_desc"] is not None:
        cve_desc = str(cve["cve_desc"]).strip()
    else:
        cve_desc = ""
    if cve["repair_time"] is not None:
        repair_time = str(cve["repair_time"]).strip()
    else:
        repair_time = ""
    update_time = str(time.strftime("%Y-%m-%d %H:%M:%S", time.localtime()))
    # State 0 means new, 1 means modified
    if cve_status in (0, 3, 4, 5, 6, 7, 8):
        cve_statux = 0
    elif cve_status in (1, 2):
        cve_statux = 1
    else:
        cve_statux = cve_status
    try:
        listx = crawltask.crawling(url)
        if (listx[2] is None or listx[2] == "" or len(listx[2]) < 2) and \
                (listx[0] is None or listx[0] == 0) and \
                (listx[3] is None or listx[3] == "" or len(listx[3]) < 2):
            print("The data does not exist and will not be processed temporarily", listx)
            return ret_value
        if repair_time == "" and listx[3] is not None and len(listx[3]) > 2:
            if len(listx[3]) > 10:
                repair_times = listx[3][:10]
            else:
                repair_times = listx[3]
            sql = "update cve_vuln_center set repair_time = %s,update_time=%s," \
                  "cve_status=%s where cve_id=%s"
            val = (repair_times, update_time, cve_statux, cve_id)
            mysql.update(sql, val)
            mysql.dispose()
        if listx[2] is not None and len(listx[2]) > 2 and cve_desc == "":
            sql = "update cve_vuln_center set cve_desc = %s,update_time=%s," \
                  "cve_status=%s where cve_id=%s"
            val = (listx[2], update_time, cve_statux, cve_id)
            mysql.update(sql, val)
            mysql.dispose()
        if listx[0] is not None and float(listx[0]) > 0:
            if listx[16] == "v3.0":
                score_type = "v3"
            else:
                score_type = "v2"
            sql = "select id,nvd_score from cve_score where cve_id = %s and " \
                  "cve_num = %s"
            val = (cve_id, cve_num)
            nvd_result = mysql.getOne(sql, val)
            if nvd_result and len(nvd_result) > 0 and nvd_result["nvd_score"] == 0:
                sql = "update cve_score set nvd_score=%s, " \
                      "n_vector_value=%s, n_attack_vector=%s, n_access_vector=%s, " \
                      "n_attack_complexity=%s, n_access_complexity=%s, n_privilege_required=%s, " \
                      "n_user_interaction=%s, n_scope=%s, n_confidentiality=%s, n_integrity=%s, " \
                      "n_availability=%s, n_authentication=%s, update_time=%s,score_type=%s " \
                      "where id = %s"
                val = (
                    listx[0], listx[4],
                    listx[5], listx[6], listx[7], listx[8], listx[9],
                    listx[10], listx[11], listx[12], listx[13], listx[14],
                    listx[15], update_time, score_type, nvd_result["id"])
                mysql.update(sql, val)
                mysql.dispose()
        cve_template_sql = "select template_id,nvd_score,nvd_vector,cve_brief from " \
                           "cve_issue_template where (nvd_score = %s or " \
                           "nvd_vector = %s or cve_brief = %s) and cve_id = %s and cve_num = %s"
        template_val = (0, "", "", cve_id, cve_num)
        template_result = mysql.getOne(cve_template_sql, template_val)
        if template_result and len(template_result) > 0:
            flag = False
            if template_result["nvd_score"] == 0 and listx[0] and float(listx[0]) > 0:
                temp_update_sql = "update cve_issue_template set nvd_score = %s," \
                                  "update_time=%s where template_id = %s"
                val = (listx[0], update_time, template_result["template_id"])
                mysql.update(temp_update_sql, val)
                mysql.dispose()
                flag = True
            if (not template_result["nvd_vector"] or len(template_result["nvd_vector"]) < 2) \
                    and (listx[4] and len(listx[4]) > 2):
                temp_update_sql = "update cve_issue_template set nvd_vector = %s," \
                                  "update_time=%s where template_id = %s"
                val = (listx[4], update_time, template_result["template_id"])
                mysql.update(temp_update_sql, val)
                mysql.dispose()
                flag = True
            if (not template_result["cve_brief"] or len(template_result["cve_brief"]) < 2) \
                    and (listx[2] and len(listx[2]) > 2):
                temp_update_sql = "update cve_issue_template set cve_brief = %s," \
                                  "update_time=%s where template_id = %s"
                val = (listx[2], update_time, template_result["template_id"])
                mysql.update(temp_update_sql, val)
                mysql.dispose()
                flag = True
            if flag:
                sql = "update cve_vuln_center set update_time=%s," \
                      "cve_status=%s where cve_id=%s"
                val = (update_time, cve_statux, cve_id)
                mysql.update(sql, val)
        mysql.dispose()
        ret_value = 1
    except IndexError as e:
        print("Subscript out of bounds", e)
        mysql.dispose(0)
    return ret_value


def supplement_cve():
    """
    1. Find cve with missing necessary fields;
    2. Go to the CVE official website to find the CVE information;
    3. Fill in the corresponding table again;
    return None
    """
    cve_list_data = list()
    mysql = Mysql()
    cve_list = query_cve_all_data(mysql, 3)
    if cve_list is not None and len(cve_list) > 0:
        for info in cve_list:
            if info not in cve_list_data:
                cve_list_data.append(info)
    if cve_list_data is not None and len(cve_list_data) > 0:
        for cve in cve_list_data:
            print(cve)
            cve_num = str(cve["cve_num"]).strip()
            cve_version = str(cve["cve_version"]).strip()
            pack_name = str(cve["pack_name"]).strip()
            url = "https://nvd.nist.gov/vuln/detail/" + cve_num
            update_cve_vuln(url, cve, mysql)
            sql = "select * from cve_spec_error where cve_num = %s and " \
                  "cve_owner = %s and pack_name = %s"
            val = (cve_num, "src-openEuler", pack_name)
            result_spec_error = mysql.getOne(sql, val)
            if result_spec_error:
                print("过滤,修改status为6：{}".format(cve_num))
                sql = "update cve_origin_excel set cve_desc = %s, cve_status = %s " \
                      "where cve_num= %s and pack_name = %s and cve_version = %s"
                val = (result_spec_error["cve_desc"], 6, cve_num, pack_name, cve_version)
                mysql.update(sql, val)
                mysql.dispose()
    mysql.close()


def long_supplement_cve():
    """
    1. Find cve with missing necessary fields;
    2. Go to the CVE official website to find the CVE information;
    3. Fill in the corresponding table again;
    return None
    """
    cve_list_data = list()
    mysql = Mysql()
    cve_list = query_cve_all_data(mysql, 1500)
    if cve_list is not None and len(cve_list) > 0:
        for info in cve_list:
            if info not in cve_list_data:
                cve_list_data.append(info)
    if cve_list_data is not None and len(cve_list_data) > 0:
        for cve in cve_list_data:
            print(cve)
            cve_num = str(cve["cve_num"]).strip()
            cve_version = str(cve["cve_version"]).strip()
            pack_name = str(cve["pack_name"]).strip()
            url = "https://nvd.nist.gov/vuln/detail/" + cve_num
            update_cve_vuln(url, cve, mysql)
            sql = "select * from cve_spec_error where cve_num = %s and " \
                  "cve_owner = %s and pack_name = %s"
            val = (cve_num, "src-openEuler", pack_name)
            result_spec_error = mysql.getOne(sql, val)
            if result_spec_error:
                print("过滤,修改status为6：{}".format(cve_num))
                sql = "update cve_origin_excel set cve_desc = %s, cve_status = %s " \
                      "where cve_num= %s and pack_name = %s and cve_version = %s"
                val = (result_spec_error["cve_desc"], 6, cve_num, pack_name, cve_version)
                mysql.update(sql, val)
                mysql.dispose()
    mysql.close()


def query_cve_by_number(mysql, cve_num):
    """
    Find cve with missing necessary fields
    """
    center_sql = "select cve_num, pack_name, cve_version," \
                 "cve_desc,repair_time,cve_status,cve_id from " \
                 "cve_vuln_center where " \
                 "cve_num = %s order by cve_id desc"
    center_val = (cve_num)
    cve_list = mysql.getMany(center_sql, center_val)
    return cve_list


def pull_cve(cve_num):
    """
    1. Find cve with missing necessary fields;
    2. Go to the CVE official website to find the CVE information;
    3. Fill in the corresponding table again;
    return None
    """
    cve_list_data = list()
    ret_value = 0
    mysql = Mysql()
    cve_list = query_cve_by_number(mysql, cve_num)
    if cve_list and len(cve_list) > 0:
        for info in cve_list:
            if info not in cve_list_data:
                cve_list_data.append(info)
    print(cve_list_data)
    if cve_list_data and len(cve_list_data) > 0:
        for cve in cve_list_data:
            print(cve)
            cve_num = str(cve["cve_num"]).strip()
            cve_version = str(cve["cve_version"]).strip()
            pack_name = str(cve["pack_name"]).strip()
            url = "https://nvd.nist.gov/vuln/detail/" + cve_num
            ret_value = update_cve_vuln(url, cve, mysql)
            sql = "select * from cve_spec_error where cve_num = %s and " \
                  "cve_owner = %s and pack_name = %s"
            val = (cve_num, "src-openEuler", pack_name)
            result_spec_error = mysql.getOne(sql, val)
            if result_spec_error:
                print("过滤,修改status为6：{}".format(cve_num))
                sql = "update cve_origin_excel set cve_desc = %s, cve_status = %s " \
                      "where cve_num= %s and pack_name = %s and cve_version = %s"
                val = (result_spec_error["cve_desc"], 6, cve_num, pack_name, cve_version)
                mysql.update(sql, val)
                mysql.dispose()
    mysql.close()
    return ret_value
